#
# Copyright (c) 2025 Industrial Software Feature Database
#


from pyNastran.op2.op2 import OP2
from pyNastran.op2.op2_geom import OP2Geom
from pyNastran.op2 import *
import time
import logging
logger = logging.getLogger(__name__)
import os
import json


def save_as_f06(model: OP2):
    start_time = time.time()
    model.write_f06(model.f06_filename)
    end_time = time.time()
    logger.debug("wrtie data to f06 successfully in " +
                 str(end_time - start_time) + " seconds")
    print("wrtie data to f06 successfully in " +
          str(end_time - start_time) + " seconds")

def get_available_results(model: OP2):
    results = {}
    table_names = model.result_names
    for table_name in table_names:
        obj = model.get_result(table_name)
        if obj is None:
            continue
        elif isinstance(obj, dict):
            if len(obj) == 0:
                continue
            else:
                results[table_name] = obj
        else:
            results[table_name] = obj
    return results

def get_first_column(data):
    # 检查维度
    if data.ndim == 1:
        # 如果是1维数组，直接返回
        return data
    elif data.ndim == 2:
        # 如果是2维数组，取第一列并展平为一维数组
        return data[:, 0].ravel()  # 使用ravel确保它是一维的
    else:
        raise ValueError("Input data must be either 1D or 2D ndarray")

def skip_pcompts(self: OP2):
    """
    Skips over the PCOMPTS or PCOMPT table without parsing any data.
    Useful in read_mode=1 or fast scan mode.
    """

    table_name = self.op2_reader._read_table_name(rewind=False)
    self.op2_reader.read_markers([-1])

    # Skip header (7 integers)
    header_data = self.op2_reader._skip_record()

    if table_name == b'PCOMPT':
        n4_mult, n5_mult = 4 * 4, 32 * 4
    elif table_name == b'PCOMPTS':
        n4_mult, n5_mult = 3 * 4, 1 * 4
    else:
        raise NotImplementedError(f"Unsupported table name: {table_name}")

    self.op2_reader.read_3_markers([-2, 1, 0])
    self.op2_reader._skip_record()  # skip 'IPCOMPT '

    self.op2_reader.read_3_markers([-3, 1, 0])
    if self.op2_reader.get_nmarkers(1, rewind=True) != [-4]:
        self.op2_reader._skip_record()

    self.op2_reader.read_3_markers([-4, 1, 0])
    if self.op2_reader.get_nmarkers(1, rewind=True) != [0]:
        self.op2_reader._skip_record()  # skip n4 block
    else:
        self.op2_reader.read_markers([0])
        return  # no n4 data, early exit

    self.op2_reader.read_3_markers([-5, 1, 0])
    self.op2_reader._skip_record()  # skip n5 block

    self.op2_reader.read_3_markers([-6, 1, 0])
    self.op2_reader.read_markers([0])

def read_op2_with_retry(op2_filename: str) -> OP2:
    model = OP2()
    methods_dict = {b'PCOMPTS': skip_pcompts}
    model.set_additional_generalized_tables_to_read(methods_dict)

    # 尝试的不同读取方式（方法名 + 显示名）
    fallback_methods = [
        (None, "默认读取"),  # 默认方式不调用 set_xxx
        ("set_as_nasa95", "NASA95"),
        ("set_as_autodesk", "AUTODESK"),
        ("set_as_nx", "NX"),
        ("set_as_optistruct", "OptiStruct")
    ]

    for method_name, label in fallback_methods:
        try:
            # 重新实例化 OP2 每次都清理干净
            model = OP2()
            if method_name:
                getattr(model, method_name)()
            model.set_additional_generalized_tables_to_read(methods_dict)
            model.read_op2(op2_filename=op2_filename)
            print(f"成功使用 [{label}] 格式读取 OP2 文件")
            return model
        except Exception as e:
            print(f"[{label}] 格式读取失败: {e}")

    raise RuntimeError("所有已知 Nastran 格式读取均失败！")

def parse_op2(file_path):

    model = read_op2_with_retry(op2_filename=file_path)

    expected_results = {'mpc_forces', 'spc_forces', 'displacements',
                        'force.ctria3_force', 'force.cquad4_force',
                        'force.cbar_force', 'force.crod_force', 'force.cbush_force',
                        'force.celas1_force', 'grid_point_forces'}
    result_names = [name for name in model.result_names if name in expected_results]


    op2_dict = {}
    result_dict = {}

    for result_name in result_names:
        table = model.get_result(result_name)
        subcase_dict = {}
        for subcase_name in table.keys():

            subcase = table.get(subcase_name)

            if hasattr(subcase, 'get_headers'):
                try:
                    stat = subcase.get_headers() # short=short
                except Exception:
                    raise

            data_dict = {}
            data_dict["headers"] = stat

            if hasattr(subcase, "element") and subcase.element is not None:
                element_data = subcase.element
                data_dict["element"] = get_first_column(element_data)
            elif hasattr(subcase, "element_node") and subcase.element_node is not None:
                element_data = subcase.element_node
                data_dict["element"] = get_first_column(element_data)
            elif hasattr(subcase, "node_element") and subcase.node_element is not None:
                element_data = subcase.node_element
                data_dict["element"] = get_first_column(element_data)
            elif hasattr(subcase, "node_gridtype") and subcase.node_gridtype is not None:
                element_data = subcase.node_gridtype
                data_dict["element"] = get_first_column(element_data)
            else:
                raise ValueError("Invalid condition: No valid attribute found in subcase.")

            data_dict["data"] = subcase.data
            subcase_dict[subcase_name] = data_dict

        result_dict[result_name] = subcase_dict

    op2_dict["title"] = model.title
    op2_dict["result"] = result_dict
    op2_dict["version"] = model._nastran_format

    return op2_dict
